package com.example;

import com.google.gson.Gson;
import com.google.gson.JsonObject;

import java.io.BufferedReader;
import java.io.InputStreamReader;
import java.io.OutputStream;
import java.net.HttpURLConnection;
import java.net.URL;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;

public class OllamaClient {
    private static final String OLLAMA_API_URL = "http://localhost:11434/api/chat";
    private static final Gson gson = new Gson();

    public static void main(String[] args) {
        try {
            // 调用封装的方法
            String result = sendMessageAndGetContent("gemma3:1b", "fuck you？");
            System.out.println(result);
        } catch (Exception e) {
            e.printStackTrace();
            System.err.println("Error occurred: " + e.getMessage());
        }
    }

    // 新增方法：封装发送消息并获取 content 的逻辑
    public static String sendMessageAndGetContent(String model, String messageContent) throws Exception {
        // 构建消息列表
        List<Map<String, String>> messages = new ArrayList<>();
        messages.add(createMessage("user", messageContent));

        // 构建请求体
        Map<String, Object> requestBody = new HashMap<>();
        requestBody.put("model", model);
        requestBody.put("messages", messages);

        // 将请求体转换为 JSON 字符串
        String jsonBody = gson.toJson(requestBody);

        // 使用 HttpURLConnection 发送 POST 请求
        URL url = new URL(OLLAMA_API_URL);
        HttpURLConnection connection = (HttpURLConnection) url.openConnection();
        connection.setRequestMethod("POST");
        connection.setRequestProperty("Content-Type", "application/json");
        connection.setDoOutput(true);

        // 写入请求体
        try (OutputStream os = connection.getOutputStream()) {
            os.write(jsonBody.getBytes("UTF-8"));
        }

        // 读取响应并解析 content 字段
        StringBuilder contentBuilder = new StringBuilder();
        try (BufferedReader br = new BufferedReader(new InputStreamReader(connection.getInputStream(), "UTF-8"))) {
            String line;
            while ((line = br.readLine()) != null) {
                if (!line.trim().isEmpty()) {
                    JsonObject jsonResponse = gson.fromJson(line, JsonObject.class);
                    if (jsonResponse.has("message") && jsonResponse.get("message").isJsonObject()) {
                        JsonObject message = jsonResponse.getAsJsonObject("message");
                        if (message.has("content") && !message.get("content").isJsonNull()) {
                            contentBuilder.append(message.get("content").getAsString());
                        }
                    }
                }
            }
        } catch (Exception e) {
            throw new RuntimeException("Error occurred while parsing response: " + e.getMessage(), e);
        }

        // 返回最终的 content 字符串
        return contentBuilder.toString().trim();
    }

    private static Map<String, String> createMessage(String role, String content) {
        Map<String, String> message = new HashMap<>();
        message.put("role", role);
        message.put("content", content);
        return message;
    }
}